#!/usr/bin/env python

"""
chr num distribution plot
"""
import re, sys, os, time
from pylab import *
import matplotlib.pyplot as plt

try:
    import psyco
    psyco.full()
except:
    pass

remode = re.compile('chr.')

def chroms(file):
    chrs = ['chr'+`i` for i in range(1,23)]
    return chrs+['chrX', 'chrY', 'chrM']


def chrom_num(file, chr):
    f = open(file, 'r')
    num = 0
    for line in f:
        line = line.strip().split()
        if line[0] == chr:
            num += 1
    f.close()
    return num

def chroms_num(file):
    chrs = chroms(file)
    dict = {}
    for chr in chrs:
        dict[chr] = chrom_num(file, chr)
    return dict

def _chroms_num_barplot(file):
    chrkeys = chroms(file)
    chroms_numbers = chroms_num(file)
    fig = figure()
    ax = fig.add_subplot(111)
    plt.subplots_adjust(bottom = 0.26)
    y = [chroms_numbers[key] for key in chrkeys]
    x = range(len(y))
    plt.bar(x,y,0.75,facecolor='purple', alpha=0.75,label='total %d' % sum(y))
    plt.legend(loc=1)
    plt.xlabel('chromosome',fontsize=10)
    plt.ylabel('numbers',fontsize=10)
    plt.title('Chromosome number distribution of %s' % file)
    plt.ylim(0,max(y)+min(y))
    loc = [ i+1 for i in x]
    labels = [i+'_'+ `chroms_numbers[i]`+'_'+ str(chroms_numbers[i]/float(sum(y))*100)[0:5] for i in chrkeys]
    plt.xticks(loc, labels, rotation=-45,fontsize=8)
    plt.savefig('%s.png' % file,transparent=True)

def chroms_num_hbarplot(file):
    chrkeys = chroms(file)
    chroms_numbers = chroms_num(file)
    fig = figure()
    ax = fig.add_subplot(111)
    val = [chroms_numbers[key] for key in chrkeys]
    pos = [i+0.5 for i in range(len(val))]
    plt.barh(pos, val, align='center', label='total %d' % sum(val),color='purple')
    plt.legend(loc=1)
    plt.yticks(pos, chrkeys)
    plt.xlabel('numbers',fontsize=12)
    plt.ylabel('chromosomes')
    plt.title('Chromosome number distribution of %s' % file)
    for (i,j) in zip(pos,val):
        plt.text(j+2*min(val),int(i),'%d-%2.2f' % (val[int(i)], val[int(i)]/float(sum(val))*100),fontsize=8,color='red')
    plt.grid(True)
    plt.xlim(0,max(val)*1.15)
    plt.savefig('%s.png' % file,transparent=True)

if __name__ == '__main__':
    import glob
    files = glob.glob('*.bed')
    for file in files:
        chroms_num_hbarplot(file)
